# -*- coding: utf-8 -*-
# @Time    : 2024/9/9 21:31
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : atom_species.py
# @Software: PyCharm
import numpy
import scipy.constants as C
from _logging import logger

class TransitionLineData:
    I_sat = 1
    mass_kg = 1
    mu_F = 1
    Gamma = 1
    f_0 = 1


def g_J_weak_mag_field(J, S, L):
    """
    按照L-S耦合
    :return:
    """
    return 1 + (J * (J + 1) + S * (S + 1) - L * (L + 1)) / (2 * J * (J + 1))


def g_F_approximately(F, J, I, g_J):
    """
    Lande g-factor
    :param F:
    :param J:
    :param I:
    :param g_J:
    :return:
    """
    return (F * (F + 1) + J * (J + 1) - I * (I + 1)) / (2 * F * (F + 1)) * g_J


def get_Gamma(omega_0, J, J_, f):
    """
    计算自然线宽，J_表示上能级
    Spontaneous decay rate Γ (Einstein A coefficient), which is also the natural (homogenous) line width (as an angular frequency) of the emitted radiation.
    Ref: https://steck.us/alkalidata/rubidium87numbers.pdf 2021 revision Eq.4
    :return:
    """
    return (C.e * omega_0) ** 2 / (2 * numpy.pi * C.epsilon_0 * C.m_e * C.c ** 3) * (2 * J + 1) / (2 * J_ + 1) * f


class Rb87_D2(TransitionLineData):
    """
    Rb 87 D2 line
    Ref: Maximized atom number for a grating magneto-optical trap via machine-learning assisted parameter optimization
    """
    # I_sat = 3.57e-3 / 1e-2 ** 2  # W/m2
    I_sat = 4.6e-3 / 1e-2 ** 2  # W/m2
    # I_sat = numpy.pi*C.h*C.c / (3*(C.c / Rb87_D2.f_0)**3)*Gamma # 约等于 1.64e-3 / 1e-2 ** 2  # W/m2 # 张立文
    mass_kg = 85.4678 * C.atomic_mass




    J = 1 / 2
    I = 3 / 2  # Ref: 印建平《原子光学》2012年第一版 Page 8
    F = I + 1 / 2
    S = 1 / 2
    L = 0
    gF = g_F_approximately(F, J, I, g_J_weak_mag_field(J, S, L))
    f_0 = 384.2304844685e12
    # omega_0 = 2 * numpy.pi * f_0
    Gamma = 38.117e6  # Hz
    mu_F = gF * C.physical_constants['Bohr magneton'][0]
if __name__ == '__main__':
    logger.info(Rb87_D2.mu_F)
